import os
import json

import torch
from PIL import Image
import torchvision.utils
import os
from modelv2 import MobileNetV2
from torchvision import transforms
import numpy as np
import matplotlib.pyplot as plt


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

net = MobileNetV2(num_classes=10).to(device)
if torch.cuda.device_count() > 1:
    net = torch.nn.DataParallel(net)
state_dict = torch.load("./MobileNetV2.pth")
net.load_state_dict(state_dict)
img_path = "../resNet/dog.jfif"


img =Image.open(img_path)

img = img.convert("RGB") #4通道转为3通道 还有一个通道是透明度
trans = transforms.Compose([transforms.ToTensor(),transforms.Resize((224,224)),transforms.Normalize([0.5,0.5,0.5],[0.5,0.5,0.5])])
img = trans(img)

# print(img.shape)

# print(img.shape)
img = torch.reshape(img,(1,3,224,224))

# print(img)

class_to_idx = ["airplane","automobile","bird","cat","deer","dog","frog","horse","ship","truck"]

label_pre_index = torch.argmax(net(img),dim=1)

label_pre = class_to_idx[label_pre_index]

print(label_pre)